{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JqjnBbmi8BPM"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PTdfUcwp8Eru"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "grBmytrShbUE"
      },
      "source": [
        "# High-performance simulations with TFF\n",
        "\n",
        "This tutorial will describe how to setup high-performance simulations with TFF\n",
        "in a variety of common scenarios.\n",
        "\n",
        "TODO: b/134543154 - Populate the content, some of the things to cover here:\n",
        "- using GPUs in a single-machine setup,\n",
        "- multi-machine setup on GCP/GKE, with and without TPUs,\n",
        "- interfacing MapReduce-like backends,\n",
        "- current limitations and when/how they will be relaxed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iPFgLeZIsZ3Q"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/federated/tutorials/simulations\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/federated/blob/v0.84.0/docs/tutorials/simulations.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/federated/blob/v0.84.0/docs/tutorials/simulations.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/federated/docs/tutorials/simulations.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yiq_MY4LopET"
      },
      "source": [
        "## Before we begin\n",
        "\n",
        "First, make sure your notebook is connected to a backend that has the relevant\n",
        "components (including gRPC dependencies for multi-machine scenarios) compiled."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_zFenI3IPpgI"
      },
      "source": [
        "Now, let's start by loading the MNIST example from the TFF website, and\n",
        "declaring the Python function that will run a small experiment loop over\n",
        "a group of 10 clients."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ke7EyuvG0Zyn"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/bin/sh: pip: command not found\n"
          ]
        }
      ],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install --quiet --upgrade tensorflow-federated"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "2dVPgxN0MdG2"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "import time\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_federated as tff\n",
        "\n",
        "source, _ = tff.simulation.datasets.emnist.load_data()\n",
        "\n",
        "\n",
        "def map_fn(example):\n",
        "  return collections.OrderedDict(\n",
        "      x=tf.reshape(example['pixels'], [-1, 784]), y=example['label']\n",
        "  )\n",
        "\n",
        "\n",
        "def client_data(n):\n",
        "  ds = source.create_tf_dataset_for_client(source.client_ids[n])\n",
        "  return ds.repeat(10).shuffle(500).batch(20).map(map_fn)\n",
        "\n",
        "\n",
        "train_data = [client_data(n) for n in range(10)]\n",
        "element_spec = train_data[0].element_spec\n",
        "\n",
        "\n",
        "keras_model = tf.keras.models.Sequential([\n",
        "    tf.keras.layers.InputLayer(input_shape=(784,)),\n",
        "    tf.keras.layers.Dense(units=10, kernel_initializer='zeros'),\n",
        "    tf.keras.layers.Softmax(),\n",
        "])\n",
        "tff_model = tff.learning.models.functional_model_from_keras(\n",
        "    keras_model,\n",
        "    loss_fn=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "    input_spec=element_spec,\n",
        "    metrics_constructor=collections.OrderedDict(\n",
        "        accuracy=tf.keras.metrics.SparseCategoricalAccuracy\n",
        "    ),\n",
        ")\n",
        "\n",
        "\n",
        "trainer = tff.learning.algorithms.build_weighted_fed_avg(\n",
        "    tff_model,\n",
        "    client_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=0.02),\n",
        ")\n",
        "\n",
        "\n",
        "def evaluate(num_rounds=10):\n",
        "  state = trainer.initialize()\n",
        "  for _ in range(num_rounds):\n",
        "    t1 = time.time()\n",
        "    result = trainer.next(state, train_data)\n",
        "    state = result.state\n",
        "    train_metrics = result.metrics['client_work']['train']\n",
        "    t2 = time.time()\n",
        "    print(\n",
        "        'train metrics {m}, round time {t:.2f} seconds'.format(\n",
        "            m=train_metrics, t=t2 - t1\n",
        "        )\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CDHJF7EIiEy-"
      },
      "source": [
        "## Single-machine simulations\n",
        "\n",
        "Now on by default."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "-V6uCS_BMoR9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "train metrics OrderedDict([('accuracy', 0.14557613)]), round time 3.11 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.181893)]), round time 1.25 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.23374486)]), round time 1.23 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.26759258)]), round time 1.18 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.31944445)]), round time 1.14 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.37222221)]), round time 1.14 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.42685184)]), round time 1.25 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.4712963)]), round time 1.35 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.5269547)]), round time 1.31 seconds\n",
            "train metrics OrderedDict([('accuracy', 0.55833334)]), round time 1.31 seconds\n"
          ]
        }
      ],
      "source": [
        "evaluate()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bZ171NhcNa3M"
      },
      "source": [
        "## Multi-machine simulations on GCP/GKE, GPUs, TPUs, and beyond...\n",
        "\n",
        "Coming very soon."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "simulations.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
